import os,cv2
from sklearn.metrics import classification_report
import numpy as np
import torch
from torch.utils.data import DataLoader,Dataset

resize_size=(32,32)

def load_pic(data_dir:str):
    data=[]
    label=[]
    for img_name in os.listdir(data_dir):
        if img_name.endswith('.jpg'):
            img = cv2.imread(os.path.join(data_dir, img_name), cv2.IMREAD_GRAYSCALE)
            img:np.ndarray = cv2.resize(img, resize_size)
            data.append(img)
            label.append(int(img_name.split('_')[0]))

    return np.array(data,np.float32),np.array(label,np.int32)

def load_vector(data_dir:str):
    data=[]
    label=[]
    for img_name in os.listdir(data_dir):
        if img_name.endswith('.jpg'):
            img = cv2.imread(os.path.join(data_dir, img_name), cv2.IMREAD_GRAYSCALE)
            img:np.ndarray = cv2.resize(img, resize_size).flatten()
            data.append(img)
            label.append(int(img_name.split('_')[0]))

    return np.array(data,np.float32),np.array(label,np.int32)

def show_result(true_y,predict_y):
    print(classification_report(true_y,predict_y))
    acc=np.mean(true_y==predict_y)
    print(f"best accuracy is {acc}\n")

def load_data_tensor(data_dir:str):
    data=[]
    label=[]
    for img_name in os.listdir(data_dir):
        if img_name.endswith('.jpg'):
            img = cv2.imread(os.path.join(data_dir, img_name), cv2.IMREAD_GRAYSCALE)
            img:np.ndarray = cv2.resize(img, resize_size)
            data.append(img)
            label.append(int(img_name.split('_')[0]))

    return torch.tensor(np.array(data),dtype=torch.float32),torch.tensor(np.array(label),dtype=torch.long)